# conda activate OpenLabCluster
# python -m lilab.OpenLabCluster_train.a2_semiseqseq_data_prepare project_repr/ project_seqseq/
import pickle
import numpy as pd
import os
import os.path as osp
import shutil
import h5py
import numpy as np
import pickle
import argparse
# from openlabcluster.utils import auxiliaryfunctions
from lilab.openlabcluster_postprocess import auxiliaryfunctions
from lilab.openlabcluster_postprocess.s1a_clipNames_inplace_parse import parse_name
from lilab.openlabcluster_postprocess.s1_merge_3_file import get_assert_1_file


project_repr = '/mnt/liying.cibr.ac.cn_Data_Temp/multiview_9/chenxf/00_BehaviorAnalysis-seq2seq/SexAge/Day55_Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/representitive_k34_filt_perc66'
project_seqseq = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/feats32-2024-04-11'


def create_copy_project(project_seqseq, project_semiseqseq):
    """
    copy the project_seqseq to project_semiseqseq
    """
    os.makedirs(project_semiseqseq, exist_ok=True)
    subdirs = ['datasets', 'label', 'models', 'output', 'videos']
    for subdir in subdirs:
        os.makedirs(osp.join(project_semiseqseq, subdir), exist_ok=True)

    shutil.copyfile(osp.join(project_seqseq, 'videos/clipNames.txt'), osp.join(project_semiseqseq, 'videos/clipNames.txt'))
    shutil.copyfile(osp.join(project_seqseq, 'datasets/data.h5'), osp.join(project_semiseqseq, 'datasets/data.h5'))
    shutil.copyfile(osp.join(project_seqseq, 'config.yaml'), osp.join(project_semiseqseq, 'config.yaml'))


def update_data_h5(project_semiseqseq, cluster_labels):
    """
    update the data.h5 with representitive label
    """
    data_h5_out = osp.join(project_semiseqseq, 'datasets/data.h5')
    with h5py.File(data_h5_out,'r+') as hf:
        if 'label' in hf: hf.pop('label')
        if 'videos' in hf: hf.pop('videos')
        hf.create_dataset('label',data=cluster_labels)
        hf.create_dataset('videos', data=np.array([0 for i in cluster_labels]))
        [cliplen, featurelen] = hf['0'].shape
    np.save(osp.join(project_semiseqseq, 'label/label.npy'), cluster_labels)
    return cliplen, featurelen


def update_cfg(project_seqseq, project_semiseqseq, nK, cliplen, featurelen):
    config_yaml = osp.join(project_semiseqseq, 'config.yaml')
    auxiliaryfunctions.edit_config(config_yaml,
        {
            'cla_dim': [int(nK)],
            'num_class': [int(nK)],
            'class_name': ['cluster_' + str(i+1) for i in range(nK)],
            'tr_modelType': 'seq2seq',
            'feature_length': featurelen,
            'batch_size': 64,
        })
    with open(config_yaml, 'r+') as f:
        # replace 'project_seqseq' with 'project_semiseqseq' in config_yaml
        f.seek(0)
        new_content = f.read().replace(project_seqseq, project_semiseqseq)
        f.seek(0)
        f.truncate()
        f.write(new_content)
    

def main(project_repr, project_seqseq):
    assert {'config.yaml', 'datasets', 'label', 'videos'} < set(os.listdir(project_seqseq))
    project_semiseqseq = osp.dirname(project_seqseq) + '/semiseq2seq_iter0'
    predpklfile = get_assert_1_file(osp.join(project_repr, '*.clippredpkl'))
    create_copy_project(project_seqseq, project_semiseqseq)

    predpkldata = pickle.load(open(predpklfile, 'rb'))
    ind_rawclip = predpkldata['ind_rawclip']
    cluster_labels_sub = predpkldata['cluster_labels']
    clipNames_sub = predpkldata['clipNames']
    cluster_labels_sub = cluster_labels_sub - cluster_labels_sub.min() + 1 #start from 1

    clipNames_file_out = osp.join(project_semiseqseq, 'videos/clipNames.txt')
    clipNames = np.array([osp.basename(f.strip()) for f in (open(clipNames_file_out, 'r').readlines())])
    assert clipNames.dtype == clipNames_sub.dtype and np.all(clipNames[ind_rawclip] == clipNames_sub)
    cluster_labels = np.zeros(len(clipNames), dtype=int)
    cluster_labels[ind_rawclip] = cluster_labels_sub
    cliplen, featurelen = update_data_h5(project_semiseqseq, cluster_labels)

    nK = cluster_labels_sub.max()
    update_cfg(project_seqseq, project_semiseqseq, nK, cliplen, featurelen)
    print(project_semiseqseq)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("dir_representitive", type=str)
    parser.add_argument("dir_seq2seq", type=str)
    args = parser.parse_args()
    main(args.dir_representitive, args.dir_seq2seq)
